/*
* @Date: 2021/3/5
* @Author: XueChengwu <xuechengwu@erayt.com>
* @Copyright: 2015-2019 Erayt, Inc.
* @Description: If you have some questions, please contact: xuechengwu@erayt.com.
*/
function Bayes(){}
Bayes.prototype.train = function(trainMatrix,trainCategory){
  var cateMaps = new Map();
  for(var i = 0;i < trainCategory.length;i++){
    if(!cateMaps.has(trainCategory[i])){
      cateMaps.set(trainCategory[i],0);
    }
    cateMaps.set(trainCategory[i],cateMaps.get(trainCategory[i]) + 1);
  }
  cateMaps.forEach(function(value, key){
    cateMaps.set(key,cateMaps.get(key)/trainCategory.length);
  });

  var labelsTree = new Map();
  var labelsCount = new Map();

  for(var i = 0;i <trainMatrix.length;i++){
    if(!labelsTree.has(trainCategory[i])){
      labelsTree.set(trainCategory[i],new Map());
      labelsCount.set(trainCategory[i],new Map());
    }
    for(var j = 0;j < trainMatrix[0].length;j++){
      if(!(labelsTree.get(trainCategory[i]).has(j))){
        labelsTree.get(trainCategory[i]).set(j,new Map());
      }
      let tmp = labelsTree.get(trainCategory[i]).get(j).has(trainMatrix[i][j])?labelsTree.get(trainCategory[i]).get(j).get(trainMatrix[i][j]):0
      labelsTree.get(trainCategory[i]).get(j).set(trainMatrix[i][j], tmp + 1);
      let tmp2 = labelsCount.get(trainCategory[i]).has(j)?labelsCount.get(trainCategory[i]).get(j):0
      labelsCount.get(trainCategory[i]).set(j,tmp2 + 1);
    }
  }

  this.cateMaps = cateMaps;
  this.labelsTree = labelsTree;
  this.labelsCount = labelsCount;

}

Bayes.prototype.save = function (){
  let ret = {cateMaps: {}, labelsTree: {}, labelsCount: {}};
  console.log(this);
  this.cateMaps.forEach((item, key) => {
    ret.cateMaps[key] = item;
  });
  this.labelsCount.forEach((item, key) => {
    let tmp = {}
    item.forEach((vo, kk) => {
      tmp[kk] = vo;
    });
    ret.labelsCount[key] = tmp;
  });
  this.labelsTree.forEach((item, key) => {
    let tmp = {};
    item.forEach((vo1, key1) => {
      let tc = {};
      vo1.forEach((vo2, key2) => {
        tc[key2] = vo2;
      });
      tmp[key1] = tc;
    });
    ret.labelsTree[key] = tmp;
  });
  return ret;
}

Bayes.prototype.load = function (obj) {
  this.cateMaps = new Map();
  for (let key in obj.cateMaps) {
    this.cateMaps.set(key, obj.cateMaps[key]);
  }
  this.labelsCount = new Map();
  for (let key in obj.labelsCount) {
    let tmp = new Map();
    for (let key1 in obj.labelsCount[key]) {
      tmp.set(parseInt(key1), obj.labelsCount[key][key1]);
    }
    this.labelsCount.set(key, tmp);
  }
  this.labelsTree = new Map();
  for(let key in obj.labelsTree) {
    let tmp = new Map();
    for(let key1 in obj.labelsTree[key]) {
      let tmp1 = new Map();
      for (let key2 in obj.labelsTree[key][key1]) {
        tmp1.set(parseInt(key2), obj.labelsTree[key][key1][key2]);
      }
      tmp.set(parseInt(key1), tmp1);
    }
    this.labelsTree.set(key, tmp);
  }
}
Bayes.prototype.predict = function(inputX){
  var predictMap = new Map();
  let keys = this.cateMaps.keys();
  let labels = [];
  for (let key of keys) {
    labels.push(key);
  }
  var sm = 0;
  for(let label of labels){
    var Pc = this.cateMaps.get(label);
    var P_w = 1;
    for(var j = 0; j < inputX.length; j++){
      var P_w_j_fz = 0,P_w_j_fm = 0;
      for(let tt of labels){
        let num = this.labelsTree.get(tt).get(j).get(inputX[j]);
        P_w_j_fz += (num === undefined ? 0 : num);
        P_w_j_fm += this.labelsCount.get(tt).get(j);
      }
      // if (P_w_j_fz === 0) {
      // 	continue;
      // }
      P_w = P_w * P_w_j_fz / P_w_j_fm;

    }
    var P_w_c = 1;
    for(var j = 0;j < inputX.length;j++){
      let tt = this.labelsTree.get(label).get(j).has(inputX[j]) ? this.labelsTree.get(label).get(j).get(inputX[j]) : 0;
      // if (tt === 0) {
      // 	console.log('tt',label,  j, inputX[j]);
      // 	continue;
      // }
      P_w_c = P_w_c * tt/this.labelsCount.get(label).get(j);
    }
    predictMap.set(label, P_w_c * Pc);
    sm += P_w_c * Pc;
  }
  for(let label of labels){
    predictMap.set(label, predictMap.get(label)/sm);
  }
  return predictMap;
}
export default Bayes
